#include <SM4.h>


const uint8_t code S_box[256] = {
    0xD6, 0x90, 0xE9, 0xFE, 0xCC, 0xE1, 0x3D, 0xB7, 0x16, 0xB6, 0x14, 0xC2,
    0x28, 0xFB, 0x2C, 0x05, 0x2B, 0x67, 0x9A, 0x76, 0x2A, 0xBE, 0x04, 0xC3,
    0xAA, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99, 0x9C, 0x42, 0x50, 0xF4,
    0x91, 0xEF, 0x98, 0x7A, 0x33, 0x54, 0x0B, 0x43, 0xED, 0xCF, 0xAC, 0x62,
    0xE4, 0xB3, 0x1C, 0xA9, 0xC9, 0x08, 0xE8, 0x95, 0x80, 0xDF, 0x94, 0xFA,
    0x75, 0x8F, 0x3F, 0xA6, 0x47, 0x07, 0xA7, 0xFC, 0xF3, 0x73, 0x17, 0xBA,
    0x83, 0x59, 0x3C, 0x19, 0xE6, 0x85, 0x4F, 0xA8, 0x68, 0x6B, 0x81, 0xB2,
    0x71, 0x64, 0xDA, 0x8B, 0xF8, 0xEB, 0x0F, 0x4B, 0x70, 0x56, 0x9D, 0x35,
    0x1E, 0x24, 0x0E, 0x5E, 0x63, 0x58, 0xD1, 0xA2, 0x25, 0x22, 0x7C, 0x3B,
    0x01, 0x21, 0x78, 0x87, 0xD4, 0x00, 0x46, 0x57, 0x9F, 0xD3, 0x27, 0x52,
    0x4C, 0x36, 0x02, 0xE7, 0xA0, 0xC4, 0xC8, 0x9E, 0xEA, 0xBF, 0x8A, 0xD2,
    0x40, 0xC7, 0x38, 0xB5, 0xA3, 0xF7, 0xF2, 0xCE, 0xF9, 0x61, 0x15, 0xA1,
    0xE0, 0xAE, 0x5D, 0xA4, 0x9B, 0x34, 0x1A, 0x55, 0xAD, 0x93, 0x32, 0x30,
    0xF5, 0x8C, 0xB1, 0xE3, 0x1D, 0xF6, 0xE2, 0x2E, 0x82, 0x66, 0xCA, 0x60,
    0xC0, 0x29, 0x23, 0xAB, 0x0D, 0x53, 0x4E, 0x6F, 0xD5, 0xDB, 0x37, 0x45,
    0xDE, 0xFD, 0x8E, 0x2F, 0x03, 0xFF, 0x6A, 0x72, 0x6D, 0x6C, 0x5B, 0x51,
    0x8D, 0x1B, 0xAF, 0x92, 0xBB, 0xDD, 0xBC, 0x7F, 0x11, 0xD9, 0x5C, 0x41,
    0x1F, 0x10, 0x5A, 0xD8, 0x0A, 0xC1, 0x31, 0x88, 0xA5, 0xCD, 0x7B, 0xBD,
    0x2D, 0x74, 0xD0, 0x12, 0xB8, 0xE5, 0xB4, 0xB0, 0x89, 0x69, 0x97, 0x4A,
    0x0C, 0x96, 0x77, 0x7E, 0x65, 0xB9, 0xF1, 0x09, 0xC5, 0x6E, 0xC6, 0x84,
    0x18, 0xF0, 0x7D, 0xEC, 0x3A, 0xDC, 0x4D, 0x20, 0x79, 0xEE, 0x5F, 0x3E,
    0xD7, 0xCB, 0x39, 0x48
};

const uint8_t code FK[16] =
        { 0xa3,0xb1,0xba,0xc6, 0x56,0xaa,0x33,0x50, 0x67,0x7d,0x91,0x97, 0xb2,0x70,0x22,0xdc };


const uint8_t code CK[32][4] = {
        0x00,0x07,0x0E,0x15,
        0x1C,0x23,0x2A,0x31,
        0x38,0x3F,0x46,0x4D,
        0x54,0x5B,0x62,0x69,
        0x70,0x77,0x7E,0x85,
        0x8C,0x93,0x9A,0xA1,
        0xA8,0xAF,0xB6,0xBD,
        0xC4,0xCB,0xD2,0xD9,
        0xE0,0xE7,0xEE,0xF5,
        0xFC,0x03,0x0A,0x11,
        0x18,0x1F,0x26,0x2D,
        0x34,0x3B,0x42,0x49,
        0x50,0x57,0x5E,0x65,
        0x6C,0x73,0x7A,0x81,
        0x88,0x8F,0x96,0x9D,
        0xA4,0xAB,0xB2,0xB9,
        0xC0,0xC7,0xCE,0xD5,
        0xDC,0xE3,0xEA,0xF1,
        0xF8,0xFF,0x06,0x0D,
        0x14,0x1B,0x22,0x29,
        0x30,0x37,0x3E,0x45,
        0x4C,0x53,0x5A,0x61,
        0x68,0x6F,0x76,0x7D,
        0x84,0x8B,0x92,0x99,
        0xA0,0xA7,0xAE,0xB5,
        0xBC,0xC3,0xCA,0xD1,
        0xD8,0xDF,0xE6,0xED,
        0xF4,0xFB,0x02,0x09,
        0x10,0x17,0x1E,0x25,
        0x2C,0x33,0x3A,0x41,
        0x48,0x4F,0x56,0x5D,
        0x64,0x6B,0x72,0x79
    };


// dir == 1 z左移，0右移
#if SOFT_CAUCULATION_51_CORE
void sm4_shift(uint8_t src_addrh,uint8_t src_addrl,uint8_t dest_addrh,uint8_t dest_addrl,uint8_t shiftn,uint8_t dir)
{
    uint32_t xdata *srcp;
    uint32_t xdata  src_data;
    uint32_t xdata  src_data_tmp;

    srcp = (uint16_t) (src_addrl) ;

    src_data = *(srcp);
    src_data_tmp = src_data;

    srcp = (uint16_t) (dest_addrl) ;

    if(dir){
        src_data = src_data << shiftn;
        src_data_tmp = src_data_tmp >> (32-shiftn);

    } else {
        src_data = src_data >> shiftn;
        src_data_tmp = src_data_tmp << (32-shiftn);
    }

    src_data  = src_data | src_data_tmp;

    *srcp = src_data;
}


void sm4_xor(uint8_t src_addrh,uint8_t src_addrl,uint8_t add_addrh,uint8_t add_addrl,uint8_t length)
{
    uint8_t xdata *srcp;
    uint8_t xdata *addp;
    uint8_t xdata  i   ;
    uint8_t xdata  src_data;
    uint8_t xdata  add_data;
    srcp = (unsigned int) (src_addrl) ;
    addp = (unsigned int) (add_addrl) ;

    for(i=0;i<length;i++){
      src_data = *(srcp+i);
      add_data = *(addp+i);
      *(srcp+i) = src_data ^ add_data;
    }
}

#elif HD_CAUCULATION_51_CORE
void sm4_shift(uint8_t src_addrh,uint8_t src_addrl,uint8_t dest_addrh,uint8_t dest_addrl,uint8_t shiftn,uint8_t dir){
    SFR_PAGE = 0x01;
    QENP_SRC_ADR_H = src_addrh;
    QENP_SRC_ADR_L = src_addrl;
    QENP_ADD_ADR_H = dest_addrh;
    QENP_ADD_ADR_L = dest_addrl;
    QENP_DAT_LEN   = 0x04;
    QENP_CN        = (0x06<<5) | shiftn;
    // while(QENP_CN&0x80);
    SFR_PAGE = 0x00;
}
void sm4_xor(uint8_t src_addrh,uint8_t src_addrl,uint8_t add_addrh,uint8_t add_addrl,uint8_t length){
    SFR_PAGE = 0x01;
    QENP_SRC_ADR_H = src_addrh;
    QENP_SRC_ADR_L = src_addrl;
    QENP_ADD_ADR_H = add_addrh;
    QENP_ADD_ADR_L = add_addrl;
    QENP_DAT_LEN   = length;
    QENP_CN        = 0x80;
    // while(QENP_CN&0x80);
    SFR_PAGE = 0x00;
}

#endif


#define keyexp_buf_sdram_addr_start_h     0x01
#define keyexp_buf_sdram_addr_start_l     0x00

#define buffer_tmp_sdram_addr_start_h     0x01
#define buffer_tmp_sdram_addr_start_l     (keyexp_buf_sdram_addr_start_l + 0x10)

#define src_data_sdram_addr_start_h       0x01
#define src_data_sdram_addr_start_l       (buffer_tmp_sdram_addr_start_l + 0x10)

#define kpi_data_sdram_addr_start_h       0x01
#define kpi_data_sdram_addr_start_l       (src_data_sdram_addr_start_l + 0x10)


uint8_t xdata KeyExp_w[4][4] _at_ keyexp_buf_sdram_addr_start_l ;// 轮秘钥计算完成后存储该轮轮秘钥


static uint8_t xdata buffer_tmp[16] _at_ buffer_tmp_sdram_addr_start_l ;
uint8_t xdata srcbuf[4][4] _at_ src_data_sdram_addr_start_l;
static uint8_t xdata prki [4] _at_ kpi_data_sdram_addr_start_l;



void SM4_KeyRound(uint8_t *key_start_addr,uint8_t step)
{
    uint8_t i;

//    uint8_t *key_start_addrp;
    if(step == 0){
        // 1. 获取原始秘钥MK
        for(i=0;i<4;i++){
            KeyExp_w[i][0] = *(key_start_addr++);
            KeyExp_w[i][1] = *(key_start_addr++);
            KeyExp_w[i][2] = *(key_start_addr++);
            KeyExp_w[i][3] = *(key_start_addr++);
        }
        // 2. 获取秘钥参数FK
        for(i=0;i<16;i++){
            buffer_tmp[i] = FK[i];
        }

        // 3. 计算TK=FK^MK
        sm4_xor(keyexp_buf_sdram_addr_start_h,keyexp_buf_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l,16);
    }

    // 4. 计算PKi=TKi+1 ^ TKi+2 ^ TKi+3 ^ CKi
    for(i=0;i<4;i++)        prki[i]=KeyExp_w[1][i];
    sm4_xor(kpi_data_sdram_addr_start_h,kpi_data_sdram_addr_start_l,keyexp_buf_sdram_addr_start_h,keyexp_buf_sdram_addr_start_l + 8 ,4);
    sm4_xor(kpi_data_sdram_addr_start_h,kpi_data_sdram_addr_start_l,keyexp_buf_sdram_addr_start_h,keyexp_buf_sdram_addr_start_l + 12,4);
    buffer_tmp[12]=CK[step][0];buffer_tmp[13]=CK[step][1];buffer_tmp[14]=CK[step][2];buffer_tmp[15]=CK[step][3];
    sm4_xor(kpi_data_sdram_addr_start_h,kpi_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l + 12,4);

    // 5. 将PKi进行S盒变换得到SKi
    for(i=0;i<4;i++)        buffer_tmp[i] = S_box[prki[i]];

    // 6. S'K0 = SKi << 13 放在buffer_tmp[4~7]
    sm4_shift(buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+4,13,1);

    // 7. S'K1 = SKi << 23 放在buffer_tmp[8~11]
    sm4_shift(buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+8,23,1);

    // 8. RKi  = TKi ^ SKi ^ S`K0 ^ S`K1
    sm4_xor(buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+4,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l + 8,4);
    sm4_xor(buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+4,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l + 0,4);
    sm4_xor(keyexp_buf_sdram_addr_start_h,keyexp_buf_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l + 4,4);


    // 9.变换KeyExp_w位置：5234 -> 2345
    prki[0]=KeyExp_w[0][0];prki[1]=KeyExp_w[0][1];prki[2]=KeyExp_w[0][2];prki[3]=KeyExp_w[0][3];
    for(i=0;i<3;i++){
        KeyExp_w[i][0] = KeyExp_w[i+1][0];
        KeyExp_w[i][1] = KeyExp_w[i+1][1];
        KeyExp_w[i][2] = KeyExp_w[i+1][2];
        KeyExp_w[i][3] = KeyExp_w[i+1][3];
    }
    KeyExp_w[3][0] = prki[0];KeyExp_w[3][1] = prki[1];KeyExp_w[3][2] = prki[2];KeyExp_w[3][3] = prki[3];

    // 执行完成后，prki和KeyExp_w[3]中存储的就是该轮轮秘钥
}

unsigned char*  SM4_Encryption(uint8_t* sdata_start_addr,uint8_t* key_start_addr,uint8_t block_max,uint8_t sm4_mode)
{
    uint8_t i,j;


    // 4. 获取原始数据
    for(i=0;i<4;i++){
        srcbuf[i][0] = *(sdata_start_addr++);
        srcbuf[i][1] = *(sdata_start_addr++);
        srcbuf[i][2] = *(sdata_start_addr++);
        srcbuf[i][3] = *(sdata_start_addr++);
    }

    for(j=0;j<32;j++){

        // 0. 获取轮秘钥rkj
        if   (sm4_mode){    // 加密模式
            SM4_KeyRound(key_start_addr,j);
        }
        else     {
            for(i=0;i< (32-j);i++) {
            SM4_KeyRound(key_start_addr,i);
            }
        }


        // 1. 计算 Xi+1 ^ Xi+2 ^ Xi+3 ^ rki
        // sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,src_data_sdram_addr_start_h,src_data_sdram_addr_start_l + 4,4);
        for(i=0;i<4;i++)        prki[i]      = srcbuf[0][i];
        for(i=0;i<4;i++)        srcbuf[0][i] = srcbuf[1][i];
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,src_data_sdram_addr_start_h,src_data_sdram_addr_start_l + 8,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,src_data_sdram_addr_start_h,src_data_sdram_addr_start_l + 12,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,keyexp_buf_sdram_addr_start_h,keyexp_buf_sdram_addr_start_l + 12,4);

        // 2. S 盒变换得到Si
        for(i=0;i<4;i++)        srcbuf[0][i] = S_box[srcbuf[0][i]];

        // 3. T0 = SI << 2  -->buffer_tmp[0~3]
        sm4_shift(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+0,2,1);

        // 4. T1 = SI << 10 -->buffer_tmp[4~7]
        sm4_shift(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+4,10,1);

        // 5. T2 = SI << 18 -->buffer_tmp[8~11]
        sm4_shift(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+8,18,1);

        // 6. T3 = SI << 24 -->buffer_tmp[12~15]
        sm4_shift(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l+12,24,1);

        // 7. Xi+4 = Xi ^ Si ^ T0 ^ T1 ^ T2 ^ T3 -->srcbuf[0]
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l +  0,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l +  4,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l +  8,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,buffer_tmp_sdram_addr_start_h,buffer_tmp_sdram_addr_start_l + 12,4);
        sm4_xor(src_data_sdram_addr_start_h,src_data_sdram_addr_start_l,kpi_data_sdram_addr_start_h  ,kpi_data_sdram_addr_start_l       ,4);

        // 8.将SRCbuf循环左移1个字节后，Xi+4,Xi+1,Xi+2,Xi+3 --> Xi+1,Xi+2,Xi+3,Xi+4
        buffer_tmp[0]=srcbuf[0][0];buffer_tmp[1]=srcbuf[0][1];buffer_tmp[2]=srcbuf[0][2];buffer_tmp[3]=srcbuf[0][3];
        for(i=0;i<3;i++){
            srcbuf[i][0] = srcbuf[i+1][0];
            srcbuf[i][1] = srcbuf[i+1][1];
            srcbuf[i][2] = srcbuf[i+1][2];
            srcbuf[i][3] = srcbuf[i+1][3];
        }
        srcbuf[3][0]=buffer_tmp[0];srcbuf[3][1]=buffer_tmp[1];srcbuf[3][2]=buffer_tmp[2];srcbuf[3][3]=buffer_tmp[3];
    }

    // 8. 反序变换
    buffer_tmp[0]= srcbuf[0][0]  ; buffer_tmp[1] = srcbuf[0][1]  ; buffer_tmp[2] = srcbuf[0][2]  ; buffer_tmp[3] = srcbuf[0][3] ;
    srcbuf[0][0] = srcbuf[3][0]  ; srcbuf[0][1]  = srcbuf[3][1]  ; srcbuf[0][2]  = srcbuf[3][2]  ; srcbuf[0][3]  = srcbuf[3][3] ;
    srcbuf[3][0] = buffer_tmp[0] ; srcbuf[3][1]  = buffer_tmp[1] ; srcbuf[3][2]  = buffer_tmp[2] ; srcbuf[3][3]  = buffer_tmp[3];

    buffer_tmp[0]= srcbuf[1][0]  ; buffer_tmp[1] = srcbuf[1][1]  ; buffer_tmp[2] = srcbuf[1][2]  ; buffer_tmp[3] = srcbuf[1][3] ;
    srcbuf[1][0] = srcbuf[2][0]  ; srcbuf[1][1]  = srcbuf[2][1]  ; srcbuf[1][2]  = srcbuf[2][2]  ; srcbuf[1][3]  = srcbuf[2][3] ;
    srcbuf[2][0] = buffer_tmp[0] ; srcbuf[2][1]  = buffer_tmp[1] ; srcbuf[2][2]  = buffer_tmp[2] ; srcbuf[2][3]  = buffer_tmp[3];

    return &srcbuf[0][0];
}


